// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include "map_concat.h"
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux::functions {
    namespace {
        // See documentation at https://prestodb.io/docs/current/functions/map.html
        template<bool EmptyForNull, bool AllowSingleArg>
        class MapConcatFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                // Ensure all input types and output type are of the same map type.
                const TypePtr &mapType = args[0]->type();
                POLLUX_CHECK_EQ(mapType->kind(), TypeKind::MAP);
                for (const VectorPtr &arg: args) {
                    POLLUX_CHECK(mapType->kindEquals(arg->type()));
                }
                POLLUX_CHECK(mapType->kindEquals(outputType));

                const size_t numArgs = args.size();
                if constexpr (!AllowSingleArg) {
                    POLLUX_CHECK_GE(numArgs, 2);
                }

                exec::DecodedArgs decodedArgs(rows, args, context);
                vector_size_t maxSize = 0;
                for (int i = 0; i < numArgs; i++) {
                    const DecodedVector *decodedArg = decodedArgs.at(i);
                    const MapVector *inputMap = decodedArg->base()->as<MapVector>();
                    const vector_size_t *rawSizes = inputMap->rawSizes();
                    rows.applyToSelected([&](vector_size_t row) {
                        if (EmptyForNull && decodedArg->is_null_at(row)) {
                            return;
                        }
                        maxSize += rawSizes[decodedArg->index(row)];
                    });
                }

                const TypePtr &keyType = outputType->as_map().keyType();
                const TypePtr &valueType = outputType->as_map().valueType();

                memory::MemoryPool *pool = context.pool();
                auto combinedKeys = BaseVector::create(keyType, maxSize, pool);
                auto combinedValues = BaseVector::create(valueType, maxSize, pool);

                // Initialize offsets and sizes to 0 so that canonicalize() will
                // work also for sparse 'rows'.
                BufferPtr offsets = allocateOffsets(rows.end(), pool);
                int *rawOffsets = offsets->asMutable<vector_size_t>();

                BufferPtr sizes = allocateSizes(rows.end(), pool);
                int *rawSizes = sizes->asMutable<vector_size_t>();

                vector_size_t offset = 0;
                rows.applyToSelected([&](vector_size_t row) {
                    rawOffsets[row] = offset;
                    // Reuse the last offset and size if null key must create empty map
                    for (int i = 0; i < numArgs; i++) {
                        const DecodedVector *decodedArg = decodedArgs.at(i);
                        if (EmptyForNull && decodedArg->is_null_at(row)) {
                            continue; // Treat NULL maps as empty.
                        }
                        const MapVector *inputMap = decodedArg->base()->as<MapVector>();
                        const vector_size_t index = decodedArg->index(row);
                        const vector_size_t inputOffset = inputMap->offsetAt(index);
                        const vector_size_t inputSize = inputMap->sizeAt(index);
                        combinedKeys->copy(
                            inputMap->mapKeys().get(), offset, inputOffset, inputSize);
                        combinedValues->copy(
                            inputMap->mapValues().get(), offset, inputOffset, inputSize);
                        offset += inputSize;
                    }
                    rawSizes[row] = offset - rawOffsets[row];
                });

                auto combinedMap = std::make_shared<MapVector>(
                    pool,
                    outputType,
                    BufferPtr(nullptr),
                    rows.end(),
                    offsets,
                    sizes,
                    combinedKeys,
                    combinedValues);

                MapVector::canonicalize(combinedMap, true);

                combinedKeys = combinedMap->mapKeys();
                combinedValues = combinedMap->mapValues();

                // Check for duplicate keys
                SelectivityVector uniqueKeys(offset);
                vector_size_t duplicateCnt = 0;
                bool throwExceptionOnDuplicateMapKeys = false;
                if (auto *ctx = context.execCtx()->queryCtx()) {
                    throwExceptionOnDuplicateMapKeys =
                            ctx->queryConfig().throwExceptionOnDuplicateMapKeys();
                }
                rows.applyToSelected([&](vector_size_t row) {
                    const int mapOffset = rawOffsets[row];
                    const int mapSize = rawSizes[row];
                    if (duplicateCnt) {
                        rawOffsets[row] -= duplicateCnt;
                    }
                    for (vector_size_t i = 1; i < mapSize; i++) {
                        if (combinedKeys->equal_value_at(
                            combinedKeys.get(), mapOffset + i, mapOffset + i - 1)) {
                            if (throwExceptionOnDuplicateMapKeys) {
                                const auto duplicateKey = combinedKeys->wrapped_vector()->toString(
                                    combinedKeys->wrapped_index(mapOffset + i));
                                POLLUX_USER_FAIL("Duplicate map key {} was found.", duplicateKey);
                            }
                            duplicateCnt++;
                            // "remove" duplicate entry
                            uniqueKeys.setValid(mapOffset + i - 1, false);
                            rawSizes[row]--;
                        }
                    }
                });

                if (duplicateCnt) {
                    uniqueKeys.updateBounds();
                    const vector_size_t uniqueCount = uniqueKeys.countSelected();

                    BufferPtr uniqueIndices = allocate_indices(uniqueCount, pool);
                    vector_size_t *rawUniqueIndices =
                            uniqueIndices->asMutable<vector_size_t>();
                    vector_size_t index = 0;
                    uniqueKeys.applyToSelected(
                        [&](vector_size_t row) { rawUniqueIndices[index++] = row; });

                    VectorPtr keys =
                            BaseVector::transpose(uniqueIndices, std::move(combinedKeys));
                    VectorPtr values =
                            BaseVector::transpose(uniqueIndices, std::move(combinedValues));

                    combinedMap = std::make_shared<MapVector>(
                        pool,
                        outputType,
                        BufferPtr(nullptr),
                        rows.end(),
                        offsets,
                        sizes,
                        std::move(keys),
                        std::move(values));
                }

                context.moveOrCopyResult(combinedMap, rows, result);
            }

            static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
                // map(K,V), map(K,V), ... -> map(K,V)
                return {
                    exec::FunctionSignatureBuilder()
                    .typeVariable("K")
                    .typeVariable("V")
                    .returnType("map(K,V)")
                    .argumentType("map(K,V)")
                    .variableArity("map(K,V)")
                    .build()
                };
            }
        };
    } // namespace

    void registerMapConcatFunction(const std::string &name) {
        exec::registerVectorFunction(
            name,
            MapConcatFunction</*EmptyForNull=*/false, /*AllowSingleArg=*/false>::
            signatures(),
            std::make_unique<MapConcatFunction<
                /*EmptyForNull=*/false,
                /*AllowSingleArg=*/false> >());
    }

    void registerMapConcatAllowSingleArg(const std::string &name) {
        exec::registerVectorFunction(
            name,
            MapConcatFunction</*EmptyForNull=*/false, /*AllowSingleArg=*/true>::
            signatures(),
            std::make_unique<MapConcatFunction<
                /*EmptyForNull=*/false,
                /*AllowSingleArg=*/true> >());
    }

    void registerMapConcatEmptyNullsFunction(const std::string &name) {
        exec::registerVectorFunction(
            name,
            MapConcatFunction</*EmptyForNull=*/true, /*AllowSingleArg=*/false>::
            signatures(),
            std::make_unique<MapConcatFunction<
                /*EmptyForNull=*/true,
                /*AllowSingleArg=*/false> >(),
            exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build());
    }
} // namespace kumo::pollux::functions
